import torch
import torch.nn as nn

class AlexNet(nn.Module):
    def __init__(self,input_width=224,input_height=224,class_num=2):
        super().__init__()
        self.class_num = class_num
        self.input_width = input_width
        self.input_height = input_height
        self.conv = nn.Sequential(
            nn.Conv2d(3,64,kernel_size=3,stride=1,padding=1),
            nn.MaxPool2d(2),
            nn.ReLU(),
            nn.Conv2d(64,128,kernel_size=3,stride=1,padding=1),
            nn.MaxPool2d(2),
            nn.ReLU(),
            nn.Conv2d(128,256,kernel_size=3,stride=1,padding=1),
            nn.MaxPool2d(2),
            nn.ReLU()
        )
        self.fc = nn.Sequential(
            nn.Linear(28 * 28 * 256,128),
            nn.Sigmoid(),
            nn.Linear(128,2),
            nn.Sigmoid()
        )
    
    def forward(self,x):
        feature = self.conv(x)
        out = self.fc(feature.view(feature.size(0),-1))
        return out

if __name__ == "__main__":
    net = AlexNet()
    data = torch.rand((1,3,224,224))
    y = net(data)
    print(y)
